import os
import torch
import pickle
import dnnlib
import legacy

from segmentation.dataset import  Triplet
from torch.utils.data.dataloader import DataLoader
from segmentation.utils import eval_net_few_shot, init_get_representation


if __name__ == '__main__':
    save_dir = r'save/00007-images-mirror-low_shot-kimg20000-batch32-color-translation-cutout'
    folder_dir = os.path.join(save_dir, 'handmark')
    pkl_dir = r'save/00007-images-mirror-low_shot-kimg20000-batch32-color-translation-cutout/handmark/Seg_L_10s_cbTrue_[4, 8, 16, 32, 64, 128, 256].pkl'
    data_dir = r'./data/CANDI-64'
    which_repre_layers = [4, 8, 16, 32, 64, 128, 256]
    resize_repre = True
    multi_class = True
    length=64
    w_steps=1000

    assert os.path.exists(os.path.join(save_dir, f'ws_{w_steps}'))
    import random
    random.seed(28)
    samples = list(range(103))
    random.shuffle(samples)

    split = [7,1,2]
    split = [int(s) for s in split]
    split = [103*s//sum(split) for s in split]
    split[1] = 103 - split[0] - split[2]
    print(f'dataset split: {split}')

    test_loader = DataLoader(
        Triplet(
            data_dir=data_dir,
            save_dir=save_dir,
            multi_class=True,
            aug=False,
            sample=samples[split[0]+split[1]:103],
            combine=True,
            return_w=True,
            length=length,
            w_steps=w_steps
        ),
        batch_size=5, drop_last=False, shuffle=False, num_workers=2, pin_memory=True
    )

    with dnnlib.util.open_url(os.path.join(save_dir, 'network-snapshot-best.pkl')) as f:
        snapshot_data = legacy.load_network_pkl(f)
        G = snapshot_data['G_ema'].eval().requires_grad_(False).cuda()
        del snapshot_data
    get_representation = init_get_representation(G, which_repre_layers, 256, 'const')

    with torch.no_grad():
        with open(os.path.join(pkl_dir), 'rb') as f:
            net = pickle.load(f)['net'].eval().requires_grad_(False).cuda()

        dice, accuracy, iou, wdice, wacc, wiou = eval_net_few_shot(net, get_representation, test_loader, True, folder_dir, 0, resize_repre, multi_class)
    
    print(dice)
    print(accuracy)
    print(iou)
    print(wdice)
    print(wacc)
    print(wiou)
    print(iou.mean())

    with open(os.path.join(folder_dir, 'auto_shot_test.txt'), 'a') as log:
        log.write(f'test dice: [{dice}]'+'\n')
        log.write(f'test accu: [{accuracy}]'+'\n')
        log.write(f'test iou : [{iou}]'+'\n')
        log.write(f'test wdice : [{wdice.item():.5f}]'+'\n')
        log.write(f'test wacc : [{wacc.item():.5f}]'+'\n')
        log.write(f'test wiou : [{wiou.item():.5f}]'+'\n')
        log.write(f'test miou : [{iou.mean().item():.5f}]'+'\n')